import subprocess
import os
import csv

# 创建结果目录
results_dir = '../results_raw'
os.makedirs(results_dir, exist_ok=True)

# 输出文件
output_file = os.path.join(results_dir, 'pipeline_cpu_baseline.csv')

# 定义批大小范围（从2^1 到 2^19）
batch_sizes = [2**i for i in range(1, 20)]
# batch_sizes = [131072]

# 打开CSV文件写入
with open(output_file, 'w', newline='') as csvfile:
    fieldnames = [
        'batch_size',
        'execute_batch_count',
        'request_throughput_rps_p50'
    ]
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    for batch_size in batch_sizes:
        print(f"Running test with batch_size={batch_size}")

        # 构建命令（不包含 --if_pipeline=0）
        command = [
            './../../bin_host/test_cpu',
            f'--batch_size={batch_size}'
        ]

        # 执行命令
        try:
            # 启动子进程
            process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,  # 合并 stderr
                text=True,
                encoding='utf-8',
                errors='ignore'
            )

            # 读取所有输出
            stdout, _ = process.communicate()

        except Exception as e:
            print(f"Error executing command for batch_size={batch_size}: {e}")
            continue

        # 打印输出（可选）
        # print(stdout)

        # 提取最后一行包含 execute_batch 的内容
        execute_count = None
        for line in reversed(stdout.splitlines()):
            # print(line)
            if "execute_batch" in line:
                parts = line.strip().split()
                for i, part in enumerate(parts):
                    if part == "execute_batch" and i + 1 < len(parts):
                        try:
                            execute_count = int(parts[i + 1])
                            break
                        except ValueError:
                            continue
                if execute_count is not None:
                    break

        if execute_count is None:
            print(f"Warning: No 'execute_batch' count found for batch_size={batch_size}")
            continue

        # 计算 throughput
        request_throughput = execute_count * batch_size / 20

        # 写入CSV
        writer.writerow({
            'batch_size': batch_size,
            'execute_batch_count': execute_count,
            'request_throughput_rps_p50': request_throughput
        })

print(f"测试完成! 结果已保存到 {output_file}")